#!/bin/bash

# YOLO模型下载脚本
# 用于Docker构建时下载预训练模型

set -e  # 遇到错误时退出

echo "🚀 开始下载YOLO模型文件..."

# 创建模型目录
mkdir -p models
cd models

# 模型配置 - 支持多个镜像源
declare -A MODELS=(
    ["yolo11n.pt"]="https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11n.pt"
    ["yolo11s.pt"]="https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11s.pt"
)

# 备用镜像源（如果GitHub下载失败）
declare -A MIRROR_MODELS=(
    ["yolo11n.pt"]="https://gitee.com/ultralytics/assets/releases/download/v8.3.0/yolo11n.pt"
    ["yolo11s.pt"]="https://gitee.com/ultralytics/assets/releases/download/v8.3.0/yolo11s.pt"
)

# 下载函数
download_model() {
    local filename=$1
    local url=$2
    local mirror_url=$3
    local max_retries=2
    local retry_count=0

    echo "📥 下载模型: $filename"
    echo "   主URL: $url"

    # 尝试主URL
    while [ $retry_count -lt $max_retries ]; do
        if wget -O "$filename" "$url" --timeout=30 --progress=bar:force 2>&1; then
            # 验证文件大小
            local file_size=$(stat -c%s "$filename" 2>/dev/null || stat -f%z "$filename" 2>/dev/null || echo "0")
            if [ "$file_size" -gt 1000000 ]; then  # 至少1MB
                echo "✅ $filename 下载成功 (大小: $(($file_size / 1024 / 1024))MB)"
                return 0
            else
                echo "⚠️ $filename 文件大小异常，重新下载..."
                rm -f "$filename"
            fi
        else
            echo "❌ $filename 主URL下载失败"
        fi

        retry_count=$((retry_count + 1))
        if [ $retry_count -lt $max_retries ]; then
            echo "🔄 重试 $retry_count/$max_retries..."
            sleep 2
        fi
    done

    # 如果主URL失败，尝试镜像URL
    if [ -n "$mirror_url" ]; then
        echo "🔄 尝试镜像源: $mirror_url"
        retry_count=0
        while [ $retry_count -lt $max_retries ]; do
            if wget -O "$filename" "$mirror_url" --timeout=30 --progress=bar:force 2>&1; then
                local file_size=$(stat -c%s "$filename" 2>/dev/null || stat -f%z "$filename" 2>/dev/null || echo "0")
                if [ "$file_size" -gt 1000000 ]; then
                    echo "✅ $filename 镜像源下载成功 (大小: $(($file_size / 1024 / 1024))MB)"
                    return 0
                else
                    echo "⚠️ $filename 镜像源文件大小异常..."
                    rm -f "$filename"
                fi
            else
                echo "❌ $filename 镜像源下载失败"
            fi

            retry_count=$((retry_count + 1))
            if [ $retry_count -lt $max_retries ]; then
                echo "🔄 镜像源重试 $retry_count/$max_retries..."
                sleep 2
            fi
        done
    fi

    echo "❌ $filename 所有源都下载失败"
    return 1
}

# 下载所有模型
success_count=0
total_count=${#MODELS[@]}

for filename in "${!MODELS[@]}"; do
    url="${MODELS[$filename]}"
    mirror_url="${MIRROR_MODELS[$filename]}"

    # 检查文件是否已存在
    if [ -f "$filename" ]; then
        file_size=$(stat -c%s "$filename" 2>/dev/null || stat -f%z "$filename" 2>/dev/null || echo "0")
        if [ "$file_size" -gt 1000000 ]; then
            echo "✅ $filename 已存在，跳过下载 (大小: $(($file_size / 1024 / 1024))MB)"
            success_count=$((success_count + 1))
            continue
        else
            echo "⚠️ $filename 存在但大小异常，重新下载..."
            rm -f "$filename"
        fi
    fi

    if download_model "$filename" "$url" "$mirror_url"; then
        success_count=$((success_count + 1))
    fi
done

# 显示结果
echo ""
echo "📊 下载结果统计:"
echo "   成功: $success_count/$total_count"
echo "   失败: $((total_count - success_count))/$total_count"

# 列出下载的文件
echo ""
echo "📁 模型文件列表:"
ls -lh *.pt 2>/dev/null || echo "   没有找到模型文件"

# 计算总大小
total_size=$(du -sh . 2>/dev/null | cut -f1 || echo "未知")
echo "   总大小: $total_size"

# 检查是否至少有一个模型下载成功
if [ $success_count -eq 0 ]; then
    echo ""
    echo "❌ 没有模型下载成功，构建将失败"
    echo "💡 建议检查网络连接或使用镜像源"
    exit 1
elif [ $success_count -lt $total_count ]; then
    echo ""
    echo "⚠️ 部分模型下载失败，但至少有一个模型可用"
    echo "🚀 继续构建..."
else
    echo ""
    echo "🎉 所有模型下载成功！"
fi

echo "✅ 模型下载脚本执行完成"
